# -*- coding: utf-8 -*-
# @Time    : 2023/4/21 20:37
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : integrals_along_specified_traj.py
# @Software: PyCharm
"""
沿着指定轨道积分
"""
import os

import matplotlib
import matplotlib.pyplot as plt
import numpy
import pandas

matplotlib.use('tkagg')
import scipy.constants as C
from scipy.interpolate import interpn
from table_to_gdf_read_ascii import parse_faster
from _logging import logger


# THETA_EPS = 1e-7


def Ek_to_gamma(Ek_eV, mass=C.m_e):
    return 1 + Ek_eV * C.eV / (mass * C.c ** 2)


def gamma_to_beta(gamma):
    return (1 - 1 / gamma ** 2) ** .5


def Ek_to_gamma_beta(Ek_eV, mass=C.m_e):
    gamma = Ek_to_gamma(Ek_eV, mass)
    beta = gamma_to_beta(gamma)
    return gamma * beta


class Shape2D:
    """
    生成的线条均包含起止点
    """
    L_EPS = 1e-4

    @staticmethod
    def arc(start_pt, start_theta, end_theta, r):
        n_pts = int(numpy.abs(r * (end_theta - start_theta) / Shape2D.L_EPS))
        thetas = numpy.linspace(start_theta, end_theta, n_pts)
        pts = numpy.zeros((n_pts, 2))
        pts[:, 0] = numpy.cos(thetas) * r
        pts[:, 1] = numpy.sin(thetas) * r

        # for i in range(n_pts):
        #     pts[i] = numpy.array([numpy.cos(thetas[i]), numpy.sin(thetas[i])]) * r
        return pts - pts[0] + start_pt

    @staticmethod
    def line(angle_relative_to_x, start_pt, end_x):
        end_y = start_pt[1] + (end_x - start_pt[0]) * numpy.tan(angle_relative_to_x)
        L = numpy.linalg.norm(numpy.array([end_x, end_y]) - start_pt)
        n_pts = int(L / Shape2D.L_EPS)
        xs = numpy.linspace(start_pt[0], end_x, n_pts)
        ys = numpy.linspace(start_pt[1], end_y, n_pts)
        return numpy.array((xs, ys)).T


def Brho(Ek_eV, mass=C.m_e, q=C.e):
    """
    磁刚度
    :return:
    """
    return Ek_to_gamma_beta(Ek_eV) * mass * C.c / q


def do_integral(pts_xyz, Bs_along_traj):
    dls = numpy.linalg.norm(numpy.diff(pts_xyz, axis=0), axis=1)
    logger.info("轨迹全长：%.10f" % dls.sum())
    return (Bs_along_traj[:-1, 1] * dls).sum()


if __name__ == '__main__':
    df = parse_faster(r"\\DESKTOP-XTVPXNT\dipole-bend2\dipole-bend-2.op3.table", 1e-3)
    # rho0_default = 1.90986
    rho0 = 1.90986# input("请输入磁铁曲率半径(mm)，否则采用默认值（%.2f mm）")# 1.90986
    xs, ys, zs = [df[col].unique() for col in ['X', 'Y', 'Z']]
    xs_magnet_start = numpy.linspace(-40e-3, 40e-3, 51)  # z = 磁体端部时的粒子横向位置
    # xs_magnet_start = numpy.linspace(-40e-3, 40e-3, 3)  # z = 磁体端部时的粒子横向位置
    xs_magnet_end = numpy.zeros(xs_magnet_start.shape)
    index_of_center_traj = numpy.where(numpy.abs(xs_magnet_start - 0) < Shape2D.L_EPS / 10)[0][0]
    curves = []
    int_Bs = []
    i = 0
    designed_deflection_angle = numpy.deg2rad(30)
    ref_integrad_angle = numpy.deg2rad(40)

    R0 = 1909.86e-3  # 参考粒子的曲率半径
    for x in xs_magnet_start:
        # integral_limits_z = [-150e-3, 150e-3]
        R = R0 +x
        arc = numpy.vstack(
        (Shape2D.arc(numpy.array([0, x]), numpy.deg2rad(90),
                     numpy.deg2rad(90) + ref_integrad_angle / 2,
                     R)[::-1],
         Shape2D.arc(numpy.array([0, x]), numpy.deg2rad(90),
                     numpy.deg2rad(90) - ref_integrad_angle / 2,
                     R)))
        pts = arc
        # pts = numpy.vstack((
        #     Shape2D.line(numpy.deg2rad(180) + designed_deflection_angle / 2, arc[0], integral_limits_z[0])[-2::-1],
        #     arc,
        #     Shape2D.line(-designed_deflection_angle / 2, arc[-1], integral_limits_z[1])[1:]
        # ))
        # plt.figure()
        # plt.plot(*pts.T, '.')
        # plt.gca().set_aspect('equal')

        pts_xyz = numpy.zeros((pts.shape[0], 3))
        pts_xyz[:, [2, 0]] = pts

        Bs = interpn((xs, ys, zs), df[['BX', 'BY', 'BZ']].values.reshape((xs.shape[0], ys.shape[0], zs.shape[0], 3)),
                     pts_xyz, bounds_error=False,  # fill_value=0.
                     )
        # plt.figure()
        # plt.plot(pts_xyz[:, 2], Bs[:, 1])
        int_B = do_integral(pts_xyz, Bs)
        int_Bs.append(int_B)
        xs_magnet_end[i] = arc[-1, 1]
        logger.info(int_B)
        curves.append(pts_xyz)
        i += 1
    relative_err = (int_Bs / int_Bs[index_of_center_traj]) - 1
    __  =  int_Bs / (rho0 +xs_magnet_start )
    relative_err_of_int_B_to_rho = (__ / __[index_of_center_traj]) - 1
    outdir = '.out'
    os.makedirs(outdir, exist_ok=True)
    res_df = pandas.DataFrame(
        numpy.array((xs_magnet_start, xs_magnet_end, int_Bs, relative_err,relative_err_of_int_B_to_rho)).T,
        columns=['磁铁入口处x/m', '磁铁出口处x/m', '积分场/(Tm)', '积分场相对差', '（积分场/曲率半径）的相对误差']).to_csv(os.path.join(outdir, '积分场均匀性.csv', ),
                                                                       encoding='GBK', index=False)

    plt.figure()
    for traj in curves:
        plt.plot(*traj[:, [2, 0]].T,  # '.'
                 )
        plt.gca().set_aspect('equal')
    plt.figure()
    plt.plot(xs_magnet_start, relative_err)

    plt.figure()
    plt.plot(xs_magnet_start, int_Bs)
